#!/usr/bin/env python

import rospy, time, math
import pandas as pd
import numpy as np
from std_msgs.msg import Int32, Float64
from geometry_msgs.msg import PoseStamped, Twist, TwistStamped
from mavros_msgs.msg import AttitudeTarget, PositionTarget, State
from mavros_msgs.srv import CommandBool, SetMode
import tf.transformations as tf
from io import StringIO

class ROS_INIT():
    def ros_init(self):
        rospy.init_node('uav_formation', anonymous=True)
        self.rate = rospy.Rate(self.h)   # Hz

        # 发布指令
        self.uav_target_pose_pub = [None]*self.uav_num
        for nm in range(self.uav_num):
            self.uav_target_pose_pub[nm] = rospy.Publisher('/uav'+str(nm+1)+'/mavros/setpoint_raw/local', PositionTarget, queue_size=1)
        
        # 订阅各无人机位置角度信息
        self.pose_sub = [None]*self.uav_num
        for nm in range(self.uav_num):
            self.pose_sub[nm] = rospy.Subscriber('/uav'+str(nm+1)+'/mavros/local_position/pose', PoseStamped, self.pose_callback,queue_size=1,callback_args=nm+1)
        

        # 订阅px4信息
        rospy.Subscriber('/uav'+str(self.uav_id)+'/mavros/state', State, self.uav_px4state_cb, queue_size=1)
        
        # 上锁与解锁服务
        self.uav_armService = rospy.ServiceProxy('/uav'+str(self.uav_id)+'/mavros/cmd/arming', CommandBool)

        # 订阅地面站指令
        rospy.Subscriber('/command', Int32, self.command_callback)

    def pose_callback(self,msg,nm):
        k = nm-1
        self.uav_pose[k][0] = msg.pose.position.x
        self.uav_pose[k][1] = msg.pose.position.y
        self.uav_pose[k][2] = msg.pose.position.z
        local_angles = tf.euler_from_quaternion([msg.pose.orientation.x,msg.pose.orientation.y,msg.pose.orientation.z,msg.pose.orientation.w])
        self.uav_pose[k][3] = local_angles[0]      # 滚转
        self.uav_pose[k][4] = local_angles[1]      # 俯仰
        self.uav_pose[k][5] = local_angles[2]      # 偏航

    def uav_px4state_cb(self,msg):
        self.px4_state = msg

    def command_callback(self, msg):
        self.command = msg.data
        print('uav',self.uav_id,'Receive command: ', self.command)
        if self.command == 0:
            self.state = 'INIT'
            print('state: ', self.state)
        elif self.command == 1:
            self.state = 'ARM'
            print('state: ', self.state)
        elif self.command == 2:
            self.state = 'TAKEOFF'
            print('state: ', self.state)
        elif self.command == 3:
            self.state = 'LAND'
            print('state: ', self.state)
        elif self.command == 4:
            self.state = 'RESET'
            print('state: ', self.state)
        elif self.command == 5:
            self.state = 'START'
            print('state: ', self.state)
        elif self.command == 6:
            self.state = 'UGV_RESET'
            print('state: ', self.state)
        elif self.command == 7:
            self.state = 'WAITING'
            print('state: ', self.state)
        elif self.command == 9:  # MISSION状态指令
            self.state = 'MISSION'
            print('state: ', self.state)
        else:
            print('state: ', self.state, ' 保持上一个状态')

class Controller(ROS_INIT):
    def __init__(self):
        ###################### 读取系统参数 #############################
        self.uav_id = rospy.get_param('uav_id', 1)    # 本机id
        print("本机是 {} 号飞机".format(self.uav_id))
        self.uav_num = rospy.get_param('uav_num', 1)    # 无人机数量
        print("集结的无人机数量是 {} ".format(self.uav_num))

        ################## 轨迹相关参数 ########################
        self.traj_path = rospy.get_param('traj_path', 'uav1的数据.txt')  # 轨迹文件路径
        self.traj_data = None  # 存储轨迹数据(时间, x, y, z, vx, vy) 不含vz
        self.mission_start_time = 0  # 任务开始时间戳
        self.mission_alt = rospy.get_param('takeoff_alt',5.0) # 任务时固定高度（二维运动）

        ################## 原有INIT参数 ############
        self.h = 20  # hz
        self.uav_pose = [[0,0,0,0,0,0] for i in range(self.uav_num)]        # 无人机位姿
        self.tar_pos = [0,0,0]                                              # 目标位置    
        self.uav_vel = [[0,0,0] for i in range(self.uav_num)]               # 无人机速度
        self.px4_state = State
        self.auto_arm = True   # 自动解锁
        self.fly_alt = [1.0, 0.8]
        self.v_max = 15  # 最大xy速度
        self.i_num = 0
        self.pre_pos = [0,0]

        ################# ROS init ########################        
        self.command = 0     
        self.state = None   
        self.ros_init()
        self.load_trajectory()  # 加载轨迹

    ################## 轨迹加载与解析 ########################
    def load_trajectory(self):
        """加载轨迹文件，解析为位置(x,y)和速度(vx,vy)，z固定为任务高度"""
        try:
            # 读取并处理文件
            with open(self.traj_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            lines[0] = lines[0].replace('，', ',')  # 替换中文逗号
            cleaned_data = StringIO(''.join(lines))
            df = pd.read_csv(cleaned_data)
            df.columns = df.columns.str.strip()

            # 提取球坐标参数
            time_stamps = df['Time'].to_numpy()
            r = df['斜距'].to_numpy()
            elev = df['高低角'].to_numpy()  # 高低角(弧度)
            azim = df['方位角'].to_numpy()  # 方位角(弧度)
            azim_dot = df['方位角变化率'].to_numpy()  # 方位角变化率

            # 计算斜距变化率和高低角变化率（数值微分）
            dt = np.mean(np.diff(time_stamps)) if len(time_stamps) > 1 else 0.05
            r_dot = np.gradient(r, dt)  # 斜距变化率
            elev_dot = np.gradient(elev, dt)  # 高低角变化率

            # 球坐标转笛卡尔坐标（仅x,y，z固定）
            x = r * np.cos(elev) * np.cos(azim)
            y = r * np.cos(elev) * np.sin(azim)
            
            print("\n====================================")
            # 计算xy速度（使用用户公式，忽略vz）
            vx = (r_dot * np.cos(elev) * np.cos(azim) 
                  - r * elev_dot * np.sin(elev) * np.cos(azim) 
                  - r * azim_dot * np.cos(elev) * np.sin(azim))
            vy = (r_dot * np.cos(elev) * np.sin(azim) 
                  - r * elev_dot * np.sin(elev) * np.sin(azim) 
                  + r * azim_dot * np.cos(elev) * np.cos(azim))
            
            # 打印前10个点和后10个点（包含位置和速度）
            print("\n========== 轨迹起始10个点 ==========")
            for i in range(min(10, len(x))):
                print(f"点{i+1}: 时间={time_stamps[i]:.3f}s, 位置(x={x[i]:.2f}m, y={y[i]:.2f}m), 速度(vx={vx[i]:.2f}m/s, vy={vy[i]:.2f}m/s)")

            print("\n========== 轨迹末尾10个点 ==========")
            for i in range(max(0, len(x)-10), len(x)):
                print(f"点{i+1}: 时间={time_stamps[i]:.3f}s, 位置(x={x[i]:.2f}m, y={y[i]:.2f}m), 速度(vx={vx[i]:.2f}m/s, vy={vy[i]:.2f}m/s)")
            print("\n====================================")

            # 存储轨迹数据：时间戳, x, y, vx, vy（z固定为mission_alt）
            self.traj_data = np.column_stack((time_stamps, x, y, vx, vy))
            print(f"轨迹加载完成，共 {len(self.traj_data)} 个点，固定高度: {self.mission_alt}m")

        except Exception as e:
            print(f"轨迹加载失败: {str(e)}")
            self.traj_data = None

    ################## MISSION状态控制函数（核心修正） ########################
    def mission_control(self):
        id = self.uav_id - 1
        current_z = self.uav_pose[id][2]

        # 1. 检查是否已到达任务高度（未到达则先定高）
        if abs(current_z - self.mission_alt) > 0.2:
            print(f"任务准备：上升至固定高度 {self.mission_alt}m（当前 {current_z:.2f}m）")
            self.pubVxyPzCmd([0, 0], self.mission_alt, id)  # 零速度+定高
            self.mission_start_time = rospy.get_time()# 到达时间
            return  #从函数中退出并返回值         
        # 2. 计算任务已执行时间
        elapsed_time = rospy.get_time() - self.mission_start_time+0.05 # 增加0.05s偏移，对齐第一个数据
        setpoint = self.get_traj_setpoint(elapsed_time)

        if setpoint is None:
            print(f"MISSION等待轨迹数据... (已执行 {elapsed_time:.2f}s)")
            return

        # 4. 解析并限制速度
        x, y, vx, vy = setpoint
        vx = np.clip(vx, -self.v_max, self.v_max)
        vy = np.clip(vy, -self.v_max, self.v_max)

        # 5. 使用原有pubVxyPzCmd发布指令（xy速度+z位置）
        self.pubVxyPzCmd([vx, vy], self.mission_alt, id)
        print(f"MISSION: 时间={elapsed_time:.2f}s, 目标位置=({x:.2f},{y:.2f}), 速度=({vx:.2f},{vy:.2f})")

    def get_traj_setpoint(self, elapsed_time):
        """根据任务时间获取xy位置和速度指令"""
        if self.traj_data is None:
            return None

        # 查找当前时间对应的轨迹段
        valid_idx = np.where(self.traj_data[:, 0] <= elapsed_time)[0]
        if len(valid_idx) == 0:
            return None  # 轨迹未开始
        if len(valid_idx) == len(self.traj_data):
            last_point = self.traj_data[-1]
            return np.array([last_point[1], last_point[2], 0.0, 0.0])  # 位置保持最后一点,速度设为0

        # 线性插值
        idx = valid_idx[-1]
        t0, x0, y0, vx0, vy0 = self.traj_data[idx]
        t1, x1, y1, vx1, vy1 = self.traj_data[idx+1]
        ratio = (elapsed_time - t0) / (t1 - t0)

        x = x0 + ratio * (x1 - x0)
        y = y0 + ratio * (y1 - y0)
        vx = vx0 + ratio * (vx1 - vx0)
        vy = vy0 + ratio * (vy1 - vy0)
        return np.array([x, y, vx, vy])

    def main(self):
        start_time = rospy.get_time()
        takeoff_pos = [0,0,0]
        land_pos = [0,0,0]
        id = self.uav_id-1
        
        while not rospy.is_shutdown():
            t = rospy.get_time() - start_time
            if self.state == 'WAITING':         
                start_time = rospy.get_time()
            
            elif self.state == 'ARM':       
                if((not self.px4_state.armed) and self.auto_arm):  
                    if self.uav_arm():
                        takeoff_pos = [self.uav_pose[id][0], self.uav_pose[id][1], self.fly_alt[0]]
                        self.state = 'WAITING'
            
            elif self.state == 'TAKEOFF':   
                self.pubPxyPzCmd(takeoff_pos, id)
                if self.dis(takeoff_pos,self.uav_pose[id]) < 0.1 and abs(takeoff_pos[2] - self.uav_pose[id][2]) < 0.1:
                    print('uav',id+1,'已完成起飞!')
                    self.state = 'WAITING'    
            
            elif self.state == 'RESET':
                init_pos = self.Formation_init()
                reset_pos = [init_pos[id*2],init_pos[id*2+1],self.fly_alt[0]]
                self.pubPxyPzCmd(reset_pos, id)
                if self.dis(reset_pos,self.uav_pose[id]) < 0.1 and abs(reset_pos[2] - self.uav_pose[id][2]) < 0.1:
                    print('uav',id+1,'已完成初始化!')
                    self.state = 'WAITING'   
            
            elif self.state == 'START':
                self.Control()
            
            elif self.state == 'LAND':
                print('uav',id+1,'收到降落指令')
                land_pos = [self.uav_pose[id][0], self.uav_pose[id][1], 0]
                self.state = 'LAND_2'
            
            elif self.state == 'LAND_2':
                if(self.uav_pose[id][2] > 0.25):
                    uav_VxyVz_sp = [0.0, 0.0, -5.0]
                    self.pubPxyVzCmd(land_pos, uav_VxyVz_sp, id)
                else:
                    uav1_VxyVz_sp = [0.0, 0.0, -0.8]
                    self.pubPxyVzCmd(land_pos, uav1_VxyVz_sp, id)
                if self.dis(land_pos,self.uav_pose[id]) < 0.1 and abs(land_pos[2] - self.uav_pose[id][2]) < 0.1:
                    print('uav',id+1,'已完成降落!')
                    self.state = 'DISARM'    
            
            elif self.state == 'DISARM':
                if self.uav_disarm():
                    break           
            ################## MISSION状态调用 ########################
            elif self.state == 'MISSION':
                self.mission_control()  # 确保调用定义的函数
           
            self.rate.sleep()

    def uav_arm(self):
        if self.uav_armService(True):
            print('uav'+str(self.uav_id)+'--已解锁！')
            return True
        else:
            print('uav'+str(self.uav_id)+'--解锁失败！')
            return False

    def uav_disarm(self):
        if self.uav_armService(False):
            print('uav'+str(self.uav_id)+'--已上锁！')
            return True
        else:
            print('uav'+str(self.uav_id)+'--上锁失败！')
            return False
        
    def pubPxyPzCmd(self, cmd_p, id):
        msg = PositionTarget()
        msg.header.stamp = rospy.Time.now()
        msg.coordinate_frame = 1    
        msg.type_mask = 0b110111111000  
        msg.position.x = cmd_p[0]   
        msg.position.y = cmd_p[1]   
        msg.position.z = cmd_p[2]   
        self.uav_target_pose_pub[id].publish(msg)

    def pubPxyVzCmd(self, cmd_p, cmd_v, id):
        msg = PositionTarget()
        msg.header.stamp = rospy.Time.now()
        msg.coordinate_frame = 1    
        msg.type_mask = 0b110111011100  
        msg.position.x = cmd_p[0]   
        msg.position.y = cmd_p[1]   
        msg.velocity.z = cmd_v[2]   
        self.uav_target_pose_pub[id].publish(msg)

    def pubVxyPzCmd(self, cmd_v, pz, id):
        msg = PositionTarget()
        msg.header.stamp = rospy.Time.now()
        msg.coordinate_frame = 1    
        msg.type_mask = 0b110111000011  
        msg.velocity.x = cmd_v[0]   
        msg.velocity.y = cmd_v[1]   
        msg.velocity.z = 0          # 不控制z速度
        msg.position.z = pz         # 固定z位置
        self.uav_target_pose_pub[id].publish(msg)

    def pubVxyzCmd(self, cmd_v, id):
        msg = PositionTarget()
        msg.header.stamp = rospy.Time.now()
        msg.coordinate_frame = 1    
        msg.type_mask = 0b110111000111  
        msg.velocity.x = cmd_v[0]   
        msg.velocity.y = cmd_v[1]   
        msg.velocity.z = 0          
        self.uav_target_pose_pub[id].publish(msg)

    def Control(self):
        id = self.uav_id-1
        k = 1.4
        cmd_vel = [0,0,0]
        vxt = (self.tar_pos[0]-self.pre_pos[0])*self.h
        vyt = (self.tar_pos[1]-self.pre_pos[1])*self.h
        cmd_vel[0] = vxt + k*(self.tar_pos[0] - self.uav_pose[id][0])
        cmd_vel[1] = vyt + k*(self.tar_pos[1] - self.uav_pose[id][1])
        print(self.dis(self.tar_pos,self.uav_pose[id]))
        self.pubVxyzCmd(cmd_vel, self.uav_id-1)
        self.pre_pos = [self.tar_pos[0],self.tar_pos[1]]

    def dis(self,a,b):
        return math.sqrt((a[0]-b[0])** 2 + (a[1]-b[1])**2)
    
    def Formation_init(self):
        radius = 10.0
        angles = [2 * math.pi * i / self.uav_num for i in range(self.uav_num)]
        positions = []
        for angle in angles:
            x = radius * math.cos(angle)
            y = radius * math.sin(angle)
            positions.extend([x, y])
        return positions
               
if __name__ == "__main__":
    formation = Controller()
    try:
        formation.main()
    except rospy.ROSInterruptException:
        pass